/********************************************************************
 *
 *  dpmm_seq_draw.cpp
 *
 *  The C++ mex file for implementing generalized DPMM sequential
 *  sampling.
 *
 *  Created by Dahua Lin, on Nov 17, 2010
 *
 ********************************************************************/

#include <mex.h>

#include <vector>
#include <valarray>

#include <math.h>
#include "psample_base.h"

using namespace std;
using namespace ddp;


template<typename T>
void addvalues(int n, const T *src, vector<T>& v)
{
    for (int i = 0; i < n; ++i)
    {
        v.push_back(src[i]);
    }
}

mxArray* to_mx(const valarray<int>& a)
{
    int n = (int)a.size();
    mxArray *mx = mxCreateNumericMatrix(1, n, mxINT32_CLASS, mxREAL);    
    int *dst = (int*)mxGetData(mx);
    for (int i = 0; i < n; ++i)
    {
        dst[i] = a[i];
    }
    return mx;
}

mxArray* to_mx(const valarray<double>& a)
{
    int n = (int)a.size();
    mxArray *mx = mxCreateDoubleMatrix(1, n, mxREAL);    
    double *dst = mxGetPr(mx);
    for (int i = 0; i < n; ++i)
    {
        dst[i] = a[i];
    }
    return mx;
}

void mprint_varr(const valarray<int>& a)
{
    int n = (int)a.size();
    for (int i = 0; i < n; ++i)
    {
        mexPrintf("%d ", a[i]);
    }
    mexPrintf("\n");
}


/**
 * A lightweight wrapper of the DPMM system configuration
 */
class DPMM_Sys
{
public:
    DPMM_Sys(const mxArray *mx) : _mx(mx)
    {
        mxLogP0 = mxGetField(mx, 0, "logp0");
        mxPriW = mxGetField(mx, 0, "priw");
        mxPriQ = mxGetField(mx, 0, "priq");
        
        _nx = (int)mxGetScalar(mxGetField(mx, 0, "nx"));
        _m0 = (int)mxGetScalar(mxGetField(mx, 0, "m0"));
    }
    
public:
    
    int nx() const
    {
        return _nx;
    }
    
    int m0() const
    {
        return _m0;
    }
    
public:
     
    const double *logp0() const
    {
        return mxGetPr(mxLogP0);
    }
    
    const double *priw() const
    {
        return _m0 == 0 ? 0 : mxGetPr(mxPriW);        
    }
    
    const double *priq() const
    {
        return _m0 == 0 ? 0 : mxGetPr(mxPriQ);
    }
    
    
private:
    const mxArray *_mx;
    
    int _nx;
    int _m0;
    
    const mxArray *mxLogP0;
    const mxArray *mxPriW;
    const mxArray *mxPriQ;
    
}; // end class DPMM_Sys



/**
 * A lightweight wrapper of the input status
 */
class DPMM_Status
{
public:
    DPMM_Status(const mxArray *mx) : _mx(mx)
    {
        mxZ = mxGetField(mx, 0, "z");
        mxId = mxGetField(mx, 0, "id");
        mxU = mxGetField(mx, 0, "u");
        mxC = mxGetField(mx, 0, "c");
        mxW = mxGetField(mx, 0, "w");
        mxTheta = mxGetField(mx, 0, "theta");
        mxLik = mxGetField(mx, 0, "lik");
        mxMaxId = mxGetField(mx, 0, "max_id");
        
        _ns = (int)mxGetNumberOfElements(mxZ);
        _m = (int)mxGetNumberOfElements(mxId);
        _pdim = (int)mxGetM(mxTheta);
        _max_id = *((const int*)mxGetData(mxMaxId));
    }
    
public:
    int ns() const
    {
        return _ns;
    }
    
    int m() const
    {
        return _m;
    }
    
    int pdim() const
    {
        return _pdim;
    }
    
    int max_id() const
    {
        return _max_id;
    }
    
public:
    const int *z() const
    {
        return (const int*)mxGetData(mxZ);
    }
    
    const int *id() const
    {
        return _m == 0 ? 0 : (const int*)mxGetData(mxId);
    }
    
    const int *u() const
    {
        return _m == 0 ? 0 : (const int*)mxGetData(mxU);
    }
    
    const int *c() const
    {
        return _m == 0 ? 0 : (const int*)mxGetData(mxC);
    }
    
    const double *w() const
    {
        return _m == 0 ? 0 : mxGetPr(mxW);
    }
    
    const mxArray *theta() const
    {
        return mxTheta;
    }
    
    const double *lik() const
    {
        return mxGetPr(mxLik);
    }            
    
private:
    const mxArray *_mx;
    
    int _ns;
    int _m;
    int _pdim;
    int _max_id;
    
    const mxArray *mxZ;
    const mxArray *mxId;
    const mxArray *mxU;
    const mxArray *mxC;
    const mxArray *mxW;
    const mxArray *mxTheta;
    const mxArray *mxLik;
    const mxArray *mxMaxId;
    
}; // end class DPMM_Status




/**
 * The class to represent a DPMM sample and
 * supports a Gibbs sampling upon it
 */
class DPMM_Sample
{
public:
    DPMM_Sample(const DPMM_Sys& sys, const DPMM_Status& status, double lalpha)
    : _nx(sys.nx()), _m0(sys.m0()), _m(status.m())
    , _logp0(sys.logp0()), _priw(sys.priw()), _priq(sys.priq()), _log_alpha(lalpha)
    , _z(status.z(), status.ns())
    , _c(status.c(), status.m())
    , _w(status.w(), status.m())
    , _lik(status.lik())
    , _ws_p(_m+1), _ws_F(_m+1)
    {                        
    }
              
public:
    int get_count(int source) 
    {
        return _c[source - 1];
    }
    

public:      
    
    /**
     * Do updating on the i-th observation
     *
     * @param i   the index of the observation
     * @param rv  a uniformly distributed random value
     *
     * @return whether need to come back to do a sampling
     */
    int update(int i, double rv)
    {
        bool draw_new = false;
        
        if (_m == 0)
        {
            return 0;
        }
        else
        {
            double *p = new double[_m+1];
            double *F = new double[_m+1];
            
            get_posterior_p(i, &(_ws_p[0]));
            int k = dsample(_m+1, &(_ws_p[0]), rv, &(_ws_F[0]));
            
            delete[] p;
            delete[] F;
            
            if (k > _m0 || (k > 0 && get_count(k) > 0))
            {
                assign(i, k);
            }
            
            return k;
        }      
    }
     
    
    void get_posterior_p(int i, double *p) const
    {
        double *lpri = new double[_m+1];
        double *llik = new double[_m+1];
        
        lpri[0] = _log_alpha;
        for (int k = 0; k < _m; ++k) lpri[k+1] = log(_w[k]); 
                
        llik[0] = _logp0[i];
        for (int k = 0; k < _m; ++k) 
        {
            llik[k+1] = _lik[i * _m + k];
        }
        
        calc_posterior(_m+1, lpri, llik, p);                                
        
        delete[] lpri;
        delete[] llik;
    }
        
    void assign(int i, int source)  
    {
        if (_z[i] != source)
        {
            int s0 = _z[i];            
            int k0 = s0 - 1;
            int k = source - 1;
            
            _z[i] = source;
            
            // decrease the old bin
            if (k0 >= 0)
            {
                -- _c[k0];
                _w[k0] -= 1;
                
                if (_c[k0] == 0)  // reset the weight
                {
                    _w[k0] = (s0 <= _m0) ? _priw[k0] * _priq[k0] : 0;
                }
            }

            // add to the new bin
            ++ _c[k];
            _w[k] += 1; 
        }
    }
    
public:
    
    mxArray *mx_z() const
    {
        return to_mx(_z);
    }
    
    mxArray *mx_c() const
    {
        return to_mx(_c);
    }
    
    mxArray *mx_w() const
    {
        return to_mx(_w);
    }                    
            
private:
            
    // number
    int _nx;
    int _m0;
    int _m;
            
    //sys  
    const double *_logp0;
    const double *_priw;
    const double *_priq;
    double _log_alpha;
       
    // status
    valarray<int> _z;   
    valarray<int> _c;
    valarray<double> _w;
    const double *_lik;
    
    // work space    
    valarray<double> _ws_p;
    valarray<double> _ws_F;
    
}; // end class DPMM_Sample



/**
 * main entry
 *
 * Input:
 *   [0]: sys:      the system configuration
 *   [1]: s0:       the original status
 *   [2]: lalpha:   the log-weight of sampling new
 *   [3]: ord:      the order of updating [one-based indices int32]
 *   [4]: rvs:      the sequence of uniformly distributed random values [double]
 *
 * Output: (updated the corresponding input)
 *   [0]: z:        the updated z
 *   [1]: c:        the updated c
 *   [2]: w:        the updated w
 *   [3]: i:        the index of observation in ord to be updated (with new particle)
 *   [4]: src:      the source to which the observation is assigned
 */
void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
    // take inputs
    
    DPMM_Sys sys(prhs[0]);    
    DPMM_Status status0(prhs[1]);
    
    const mxArray *mxLAlpha = prhs[2];
    const mxArray *mxOrd = prhs[3];
    const mxArray *mxRvs = prhs[4];
        
    double lalpha = mxGetScalar(mxLAlpha);
    int nu = mxGetNumberOfElements(mxOrd);
    const int *ord = (const int*)mxGetData(mxOrd);
    const double *rvs = mxGetPr(mxRvs);
    
    // main
    int m0 = sys.m0();
    DPMM_Sample dpmms(sys, status0, lalpha);  
    
    bool snew = false;
    int i = 0; 
    int src = 0;
    for (; i < nu && !snew; ++i)    
    {        
        src = dpmms.update(ord[i]-1, rvs[i]);
        if (src == 0 || (src <= m0 && dpmms.get_count(src) == 0)) 
            break;
    }
        
    // output
    
    plhs[0] = dpmms.mx_z();
    plhs[1] = dpmms.mx_c();
    plhs[2] = dpmms.mx_w();    
    plhs[3] = mxCreateDoubleScalar(i + 1);  
    plhs[4] = i < nu ? mxCreateDoubleScalar(src) : mxCreateDoubleMatrix(0,0,mxREAL);
    
}



